import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet18
import numpy as np

class DoubleConv(nn.Module):
    """(convolution =&gt; [BN] =&gt; ReLU) * 2"""

    def __init__(self, in_channels, out_channels, mid_channels=None):
        super().__init__()
        if not mid_channels:
            mid_channels = out_channels
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(mid_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels , in_channels // 2, kernel_size=2, stride=2)
            self.conv = DoubleConv(in_channels, out_channels)


    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = x2.size()[2] - x1.size()[2]
        diffX = x2.size()[3] - x1.size()[3]

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])
        # if you have padding issues, see
        # https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
        # https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)

class ModifiedUNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(ModifiedUNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_channels, 64)  # 48
        self.down1 = Down(64, 128)  # 24
        self.down2 = Down(128, 256)  # 12
        self.down3 = Down(256, 512)  # 12
        self.down4 = Down(512, 1024)  # 6
        factor = 2 if bilinear else 1
        self.down5 = Down(1024, 2048 // factor)  # 3
        self.up1 = Up(2048, 1024 // factor, bilinear)  # 6
        self.up2 = Up(1024, 512 // factor, bilinear)  # 12
        self.up3 = Up(512, 256 // factor, bilinear)  # 24
        self.up4 = Up(256, 128 // factor, bilinear)  # 48
        self.up5 = Up(128, 64, bilinear)  # 48
        self.regression = OutConv(64, n_classes)  # 48


        # self.inc = DoubleConv(n_channels, 16) # 48
        # self.down1 = Down(16, 32) # 24
        # self.down2 = Down(32, 64) # 12
        # self.down3 = Down(64, 128) # 6
        # factor = 2 if bilinear else 1
        # self.down4 = Down(128, 256 // factor) # 3
        # self.up1 = Up(256, 128 // factor, bilinear) # 6
        # self.up2 = Up(128, 64 // factor, bilinear)# 12
        # self.up3 = Up(64, 32 // factor, bilinear)# 24
        # self.up4 = Up(32, 16, bilinear) # 48
        # self.regression = OutConv(16, n_classes)# 48
        # self.classification = resnet18(pretrained = True)
        # in_features = self.classification.fc.in_features
        # self.classification.fc = nn.Linear(in_features, 2)

    def forward(self, x):
        e0 = self.inc(x)
        # print("e0: ", e0.shape)
        e1 = self.down1(e0)
        # print("e1: ", e1.shape)

        e2 = self.down2(e1)
        # print("e2: ", e2.shape)

        e3 = self.down3(e2)
        # print("e3: ", e3.shape)

        e4 = self.down4(e3)
        # print("e4: ", e4.shape)

        e5 = self.down5(e4)
        # print("e5: ", e5.shape)

        d1 = self.up1(e5, e4)
        # print("d1: ", d1.shape)

        d2 = self.up2(d1, e3)
        # print("d2: ", d2.shape)

        d3 = self.up3(d2, e2)
        # print("d3: ", d3.shape)

        d4 = self.up4(d3, e1)
        # print("d4: ", d4.shape)

        d5 = self.up5(d4, e0)
        # print("d5: ", d5.shape)

        regression = self.regression(d5)
        # print("regression: ", regression.shape)

        regression = regression + x
        # print("regression: ", regression.shape)
        # classification = self.classification(x + regression)
        # feat = [e0, e1, e2, e3, e4, e5, d1, d2, d3, d4, regression]
        return regression




        # e1 = self.inc(x)
        # print("e1: ", e1.shape)
        #
        # e2 = self.down1(e1)
        # print("e2: ", e2.shape)
        #
        # e3 = self.down2(e2)
        # print("e3: ", e3.shape)
        #
        # e4 = self.down3(e3)
        # print("e4: ", e4.shape)
        #
        # e5 = self.down4(e4)
        # print("e5: ", e5.shape)
        #
        # d1 = self.up1(e5, e4)
        # print("d1: ", d1.shape)
        #
        # d2 = self.up2(d1, e3)
        # print("d2: ", d2.shape)
        #
        # d3 = self.up3(d2, e2)
        # print("d3: ", d3.shape)
        #
        # d4 = self.up4(d3, e1)
        # print("d4: ", d4.shape)
        #
        # regression = self.regression(d4)
        # print("regression: ", regression.shape)
        #
        # regression = regression + x
        # # classification = self.classification(x + regression)
        # feat = [e1, e2, e3, e4, e5, d1, d2, d3, d4, regression ]
        # return regression, feat
        # return regression, classification, feat


if __name__ == "__main__":
    model = ModifiedUNet(3, 3)
    x = torch.randn(4,3,224,224)
    regression = model(x)
    # regression, classification, feat = model(x)
    # regression, feat = model(x)

    print('regression.size(): ', regression.size())
    # print('classification.size(): ', classification.size())
    # print('np.array(feat).shape: ', np.array(feat).shape)
    # print('feat[0].shape: ', feat[0].shape)
    # print('feat[1].shape: ', feat[1].shape)
    # print('feat[2].shape: ', feat[2].shape)
    # print('feat[3].shape: ', feat[3].shape)
    # print('feat[4].shape: ', feat[4].shape)
    # print('feat[5].shape: ', feat[5].shape)
    # print('feat[6].shape: ', feat[6].shape)
    # print('feat[7].shape: ', feat[7].shape)
    # print('feat[8].shape: ', feat[8].shape)
    # print('feat[9].shape: ', feat[9].shape)
    # print('feat[10].shape: ', feat[10].shape)

    '''
feat = [e1, e2, e3, e4, e5, d1, d2, d3, d4, regression ]
regression.size():  torch.Size([4, 3, 224, 224])
feat[0].shape:  torch.Size([4, 16, 224, 224])
feat[1].shape:  torch.Size([4, 32, 112, 112])
feat[2].shape:  torch.Size([4, 64, 56, 56])
feat[3].shape:  torch.Size([4, 128, 28, 28])
feat[4].shape:  torch.Size([4, 128, 14, 14])
feat[5].shape:  torch.Size([4, 64, 28, 28])
feat[6].shape:  torch.Size([4, 32, 56, 56])
feat[7].shape:  torch.Size([4, 16, 112, 112])
feat[8].shape:  torch.Size([4, 16, 224, 224])
feat[9].shape:  torch.Size([4, 3, 224, 224])
    '''
